import cv2
import time
import os
#------------------------------#
import Model  ##### 添加模型 #####    需要修改的参数
#------------------------------#

def videos_seek(start_path):
    video_list = []  # 视频列表
    def video_seek():
        os.chdir(start_path)  # 查找视频的途径
        items = os.listdir(os.curdir)  # 该路径下所有文件
        video_type = ['.mp4', '.avi']  # 添加查找视频的格式   注意加.  如'.MP4'
        for each in items:
            if os.path.splitext(each)[1].lower() in video_type:
                video_list.append(os.getcwd() + os.sep + each)
            if os.path.isdir(each):
                video_seek(each)  # 递归函数查找下一目录下的文件
                os.chdir(os.pardir)  # 返回上一层文件目录
        return video_list
    video_seek()
    return video_list
    
    

def video_save(model, video_path, viveo_name, save_path):

    capture = cv2.VideoCapture(video_path)  # 加载视频

    total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))   # 视频总帧数

    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取视频的宽度

    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取视频的高度

    fps = capture.get(cv2.CAP_PROP_FPS)  # 获取视频的帧率

    fourcc = int(capture.get(cv2.CAP_PROP_FOURCC))  # 视频的编码

    writer = cv2.VideoWriter(save_path + '\\' + viveo_name + '.mp4', fourcc, fps, (width, height))   # 视频存储的相关参数

    t1 = time.time()

    frame_count = 0

    write_fps = 0.0

    while True:

        ref, frame = capture.read()

        #按q或者当视频识别到最后一帧图像结束时退出
        if (cv2.waitKey(20) & 0xFF == ord('q')) or ref == False:
            break

        frame_count += 1

        t3 = time.time()
        #---------------------------------------------------#
        frame = model.predict(frame)  ##### 模型的预测函数 #####     需要修改的参数
        # ---------------------------------------------------#
        t4 = time.time()

        write_fps = (1. / (t4 - t3) + write_fps) / 2  # 计算视频识别是的fps

        print('write_fps = {:.2f},已完成{}/{}'.format(write_fps, frame_count, total))

        frame = cv2.putText(frame, 'write_fps = {:.2f}'.format(write_fps),
                            (int(width / 15), int(height / 10)), cv2.FONT_HERSHEY_PLAIN,
                            (int(width / 600) + int(height / 360)), (0, 255, 0), (int(width / 400) + int(height / 300)))  # 添加fps至视频中
        cv2.imshow('frame', frame)  # 显示识别窗口

        writer.write(frame)  # 视频写入

    t2 = time.time()

    total_time = t2 - t1

    print('视频总共有{}帧，识别共耗时{}秒'.format(frame_count, total_time))

    capture.release()

    writer.release()

    cv2.destroyAllWindows()

if __name__ == "__main__":
    model = Model()

    videos_path = r'C:\Users\Administrator\Desktop'  # 需要识别视频的位置

    save_path = r'C:\Users\Administrator\Desktop\video'  # 存储视频的位置

    videos = videos_seek(videos_path)

    for video_path_name in videos:

        video_path = video_path_name.split('.')

        video_name = video_path[0].split('\\')[-1]

        print('总共有{}个视频，开始对{}视频进行处理。'.format(len(videos), video_name))

        video_save(model, video_path_name, video_name, save_path)




        



        





